{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TorchSparse for MMDetection3D Plugin Demo\n",
    "You can run the cells below to run the evaluation of TorchSparse integrated MMDetection3D models. \n",
    "\n",
    "## Dependencies\n",
    "- MMDetection3D installation: Please follow the [MMDetection3D documentation](https://mmdetection3d.readthedocs.io/en/latest/get_started.html). \n",
    "- Pre-process the datasets required by MMDetection3D ([see here](https://mmdetection3d.readthedocs.io/en/latest/user_guides/dataset_prepare.html)). \n",
    "- TorchSparse installation. \n",
    "- Install TorchSparse plugin for MMDetection3D\n",
    "    1. Clone this repository\n",
    "    2. Go to `examples/mmdetection3d` and run `pip install -v -e .`\n",
    "\n",
    "## Notes\n",
    "1. For model evaluation, you need to change the data root in the original mmdetection3d's model config to be the full path of the corresponding dataset root. The default is the relative path because mmdet3d expect you to run the evaluation under their repository folder. However, to run this demo, the relative path won't work and you need to change it to the full path. \n",
    "\n",
    "# Steps\n",
    "1. Install the dependencies. \n",
    "2. Specify the base pathes and model registry. \n",
    "3. Activate the plugin: In `mmdetection3d/tools/test.py`, add `import ts_plugin` as the last import statement.  \n",
    "4. Run demo. \n",
    "5. Print the evaluation results. \n",
    "\n",
    "# Lists of Supported Models\n",
    "- SECOND\n",
    "- PV-RCNN\n",
    "- CenterPoint\n",
    "- PartA2\n",
    "\n",
    "# The Actual Part\n",
    "## Load the Weight Conversion Module\n",
    "The dimensions of TorchSparse differs from the SpConv, so the parameter dimension conversion is required to use the TorchSparse backend. The following cell loads the converter. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib.util\n",
    "import sys, os\n",
    "from pathlib import Path\n",
    "import subprocess\n",
    "\n",
    "# Define the relative path to the file\n",
    "relative_path = \"../converter.py\"\n",
    "file_path = Path().resolve() / relative_path\n",
    "\n",
    "# Add the directory containing the file to sys.path\n",
    "sys.path.append(str(file_path.parent))\n",
    "\n",
    "# Load the module\n",
    "spec = importlib.util.spec_from_file_location(\"convert_weights\", str(file_path))\n",
    "converter = importlib.util.module_from_spec(spec)\n",
    "spec.loader.exec_module(converter)\n",
    "\n",
    "converter = getattr(converter, \"convert_weights\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dummy check for whether the weight converter is successfully loaded. \n",
    "print(converter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specify the Paths and Environment Parameters\n",
    "To run this demo, you need to provide the following paths:\n",
    "1. `mmdet3d_path`: MMDetection3D installation path. We need this path to find the `test.py` evaluation script. \n",
    "2. `mmdet3d_model_base_path`: Input pretrained weight path. the pretrained weights you download from the MMDetection3D model zoo should be put under a same base folder. \n",
    "3. `torchsparse_model_base_path`: Output pretrained weight path. The converted weights for various models should be put under the same base folder as well. \n",
    "4. `mmdet3d_cfg_base_path`: MMDetection3D configuration files base path. This configuration file is required in the model conversion. Specifically, it use the original configuration file to create a model to identify the Sparse Conv modules, and convert the weights for only those modules. By default, if you installed the mmdet3d in development mode with `-e` then this should just be the `config` folder in the mmdet3d repo. \n",
    "5. Conda environment name: this demo initialize a sub-shell to execute the demo with `subprocess`. So you need to specify the name of the conda environment that you want to use to run the demo. \n",
    "\n",
    "For paths 2, 3, and 4, we expect you to organize them by having a base path and put the checkpoint/configurations files of different models under the same basepath. For example, for the input pertrained weight path, the file structure looks like: \n",
    "\n",
    "```text\n",
    "mmdet_model_base_folder/                      \n",
    "├── SECOND/                 \n",
    "│   └── SECOND_Checkpoint.pth\n",
    "├── PV-RCNN/\n",
    "│   └── PV-RCNN_Checkpoint.pth\n",
    "└── CenterPoint/\n",
    "    └── CenterPoint_Checkpoint.pth\n",
    "```\n",
    "To configure the path for SECOND demo, you need to configure the `mmdet3d_model_base_path` to the path of the folder `mmdet_model_base_folder` and in the SECOND's registry entry, set `ckpt_before` to be `SECOND/SECOND_Checkpoint.pth`. \n",
    "\n",
    "In addition to the paths, we also need you to specify:\n",
    "1. SpConv version of the original model.\n",
    "2. `cfg_options`: some modules in the model is replaced by the TorchSparse layers. When running the evaluation, you don't need to provide a new configuration file to specify the use of TorchSparse layers. You can rather use the original mmdetection3d config file but use the `cfg_options` to replace certain modules to use the TorchSparse module. Typically, only one or two modules needed to be replaced. You can see the specific usage from the exaple below. \n",
    "3. Name of the conda environment the dependencies is installed. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "env_name = \"torchsparse\"\n",
    "\n",
    "# Please complete the following base paths. \n",
    "base_paths = {\n",
    "    'mmdet3d_path': None,\n",
    "    'mmdet3d_model_base_path': None,\n",
    "    'torchsparse_model_base_path': os.path.join(os.path.abspath(''), \"converted_models\"),\n",
    "    'mmdet3d_cfg_base_path': None\n",
    "}\n",
    "\n",
    "# Specify the model specific path and registry values. \n",
    "# NOTE: ckpt_before is associated with the mmdet3d_model_base_path and ckpt_after is associated with the torchsparse_model_base_path. \n",
    "second_3d_car = {\n",
    "    'ckpt_before': 'SECOND/second_hv_secfpn_8xb6-80e_kitti-3d-car-75d9305e.pth',\n",
    "    'ckpt_after': 'SECOND/second_hv_secfpn_8xb6-80e_kitti-3d-car-75d9305e.pth',\n",
    "    'cfg_path': 'second/second_hv_secfpn_8xb6-80e_kitti-3d-car.py',\n",
    "    'v_spconv': 2,\n",
    "    'cfg_options': \"--cfg-options test_evaluator.pklfile_prefix=outputs/torchsparse/second --cfg-options model.middle_encoder.type=SparseEncoderTS\"\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'mmdet3d_path': '/home/yingqi/repo/mmdetection3d',\n",
       " 'mmdet3d_model_base_path': '/home/yingqi/repo/mmdetection3d/models',\n",
       " 'torchsparse_model_base_path': '/home/yingqi/repo/torchsparse-dev/examples/mmdetection3d/converted_models',\n",
       " 'mmdet3d_cfg_base_path': '/home/yingqi/repo/mmdetection3d/configs'}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "base_paths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function to run a single demo is defined below. Based on the configuration dictionary you provid, it convert the model weights then use the `tools/test.py` in the `mmdetection3d` repo to run the model evaluation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mmdet3d_single_demo(registry_entry, base_paths, convert=True):\n",
    "    \"\"\"Run Single Model Demo\n",
    "\n",
    "    :param registry_entry: the model demo registry. \n",
    "    :type registry_entry: dict\n",
    "    :param base_paths: the base paths. \n",
    "    :type base_paths: dict\n",
    "    :param convert: whether to convert the model. If set to false, it skip the model conversion and use the provided checkpoint to run model evaluation directly. Defaults to True.\n",
    "    :type convert: bool, optional\n",
    "    :return: return the process object that used to run the demo. \n",
    "    :rtype: CompletedProcess\n",
    "    \"\"\"\n",
    "\n",
    "    assert os.path.isdir(base_paths['mmdet3d_path']), \"Please specify the mmdet3d_path in the base_paths.\"\n",
    "    assert os.path.isdir(base_paths['mmdet3d_model_base_path']), \"Please specify the mmdet3d_model_base_path in the base_paths.\"\n",
    "    assert os.path.isdir(base_paths['torchsparse_model_base_path']), \"Please specify the torchsparse_model_base_path in the base_paths.\"\n",
    "    assert os.path.isdir(base_paths['mmdet3d_cfg_base_path']), \"Please specify the mmdet3d_cfg_base_path in the base_paths.\"\n",
    "\n",
    "    # pre-process paths\n",
    "    cfg_path = os.path.join(base_paths['mmdet3d_cfg_base_path'], registry_entry['cfg_path'])\n",
    "    test_file_path = os.path.join(base_paths['mmdet3d_path'], \"tools/test.py\")\n",
    "    mmdet3d_model_path = os.path.join(base_paths['mmdet3d_model_base_path'], registry_entry['ckpt_before'])\n",
    "    assert os.path.isdir(base_paths['torchsparse_model_base_path']), \"Please create the directory for the converted model.\"\n",
    "    torchsparse_model_path = os.path.join(base_paths['torchsparse_model_base_path'], registry_entry['ckpt_after'])\n",
    "    \n",
    "    cfg_options = registry_entry['cfg_options']\n",
    "    # convert the model\n",
    "    if convert:\n",
    "        parent_dir = os.path.dirname(torchsparse_model_path)\n",
    "        if not os.path.exists(parent_dir):\n",
    "            os.makedirs(parent_dir)\n",
    "        converter(\n",
    "            ckpt_before=mmdet3d_model_path,\n",
    "            ckpt_after=torchsparse_model_path,\n",
    "            cfg_path=cfg_path,\n",
    "            v_spconv = registry_entry['v_spconv']\n",
    "        )\n",
    "\n",
    "    command = f'bash -c \"conda activate {env_name}; python {test_file_path} {cfg_path} {torchsparse_model_path} {cfg_options} --task lidar_det\"'\n",
    "    print(command)\n",
    "    result = subprocess.run(command, capture_output=True, text=True, shell=True, executable='/bin/bash')\n",
    "    return result  # result have .stdout and .stderr attributes to get the output. \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate MMDetection3d Models\n",
    "\n",
    "### SECOND\n",
    "Run a SECOND demo. You can print the evaluation results of the model from the sub-process's `stdout` and `stderr`. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "second_results = mmdet3d_single_demo(second_3d_car, base_paths, convert=True)\n",
    "print(second_results.stderr)\n",
    "print(second_results.stdout)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Expected Output: \n",
    "\n",
    "```\n",
    "----------- AP11 Results ------------\n",
    "\n",
    "Car AP11@0.70, 0.70, 0.70:\n",
    "bbox AP11:95.2015, 89.6519, 88.0073\n",
    "bev  AP11:89.9621, 87.2725, 84.2825\n",
    "3d   AP11:88.3629, 78.2199, 76.0327\n",
    "aos  AP11:94.94, 89.08, 87.23\n",
    "Car AP11@0.70, 0.50, 0.50:\n",
    "bbox AP11:95.2015, 89.6519, 88.0073\n",
    "bev  AP11:95.3329, 89.9520, 88.7400\n",
    "3d   AP11:95.2805, 89.8595, 88.5336\n",
    "aos  AP11:94.94, 89.08, 87.23\n",
    "\n",
    "----------- AP40 Results ------------\n",
    "\n",
    "Car AP40@0.70, 0.70, 0.70:\n",
    "bbox AP40:97.4063, 92.4550, 89.2481\n",
    "bev  AP40:92.6387, 88.4049, 85.2355\n",
    "3d   AP40:90.4511, 81.3433, 76.1927\n",
    "aos  AP40:97.13, 91.81, 88.42\n",
    "Car AP40@0.70, 0.50, 0.50:\n",
    "bbox AP40:97.4063, 92.4550, 89.2481\n",
    "bev  AP40:97.5160, 94.7415, 91.7295\n",
    "3d   AP40:97.3701, 94.5687, 91.4920\n",
    "aos  AP40:97.13, 91.81, 88.42\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PV-RCNN\n",
    "Run a PV-RCNN Demo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PV-RCNN Registry\n",
    "pv_rcnn_config = {\n",
    "    \"ckpt_before\": \"PV-RCNN/pv_rcnn_8xb2-80e_kitti-3d-3class_20221117_234428-b384d22f.pth\",\n",
    "    \"ckpt_after\": \"PV-RCNN/pv_rcnn_8xb2-80e_kitti-3d-3class_20221117_234428-b384d22f.pth\",\n",
    "    \"cfg_path\": \"pv_rcnn/pv_rcnn_8xb2-80e_kitti-3d-3class.py\",\n",
    "    \"v_spconv\": 1,\n",
    "    \"cfg_options\": \"--cfg-options test_evaluator.pklfile_prefix=outputs/torchsparse/pv_rcnn --cfg-options model.middle_encoder.type=SparseEncoderTS --cfg-options model.points_encoder.type=VoxelSetAbstractionTS\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pv_rcnn_results = mmdet3d_single_demo(pv_rcnn_config, base_paths, convert=True)\n",
    "print(pv_rcnn_results.stderr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(pv_rcnn_results.stdout)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Expected Output: \n",
    "\n",
    "```\n",
    "----------- AP11 Results ------------\n",
    "\n",
    "Pedestrian AP11@0.50, 0.50, 0.50:\n",
    "bbox AP11:74.1319, 68.4703, 65.9149\n",
    "bev  AP11:68.2026, 62.7491, 57.1043\n",
    "3d   AP11:66.6080, 59.7569, 55.1617\n",
    "aos  AP11:68.98, 63.64, 60.68\n",
    "Pedestrian AP11@0.50, 0.25, 0.25:\n",
    "bbox AP11:74.1319, 68.4703, 65.9149\n",
    "bev  AP11:80.5667, 76.2589, 72.7974\n",
    "3d   AP11:80.3568, 76.0134, 72.2977\n",
    "aos  AP11:68.98, 63.64, 60.68\n",
    "Cyclist AP11@0.50, 0.50, 0.50:\n",
    "bbox AP11:89.2310, 82.0387, 77.1643\n",
    "bev  AP11:87.8058, 74.9448, 70.5274\n",
    "3d   AP11:87.2027, 73.2608, 69.6121\n",
    "aos  AP11:89.13, 81.69, 76.77\n",
    "Cyclist AP11@0.50, 0.25, 0.25:\n",
    "bbox AP11:89.2310, 82.0387, 77.1643\n",
    "bev  AP11:88.6302, 80.1792, 74.8060\n",
    "3d   AP11:88.6302, 80.1792, 74.8060\n",
    "aos  AP11:89.13, 81.69, 76.77\n",
    "Car AP11@0.70, 0.70, 0.70:\n",
    "bbox AP11:96.0265, 89.5369, 89.1852\n",
    "bev  AP11:90.1265, 88.0958, 87.6436\n",
    "3d   AP11:89.2321, 83.7058, 78.7935\n",
    "aos  AP11:95.98, 89.43, 89.03\n",
    "Car AP11@0.70, 0.50, 0.50:\n",
    "bbox AP11:96.0265, 89.5369, 89.1852\n",
    "bev  AP11:96.1496, 94.8182, 89.2712\n",
    "3d   AP11:96.0921, 89.5371, 89.2317\n",
    "aos  AP11:95.98, 89.43, 89.03\n",
    "\n",
    "Overall AP11@easy, moderate, hard:\n",
    "bbox AP11:86.4631, 80.0153, 77.4215\n",
    "bev  AP11:82.0450, 75.2632, 71.7584\n",
    "3d   AP11:81.0143, 72.2412, 67.8558\n",
    "aos  AP11:84.70, 78.25, 75.49\n",
    "\n",
    "----------- AP40 Results ------------\n",
    "\n",
    "Pedestrian AP40@0.50, 0.50, 0.50:\n",
    "bbox AP40:75.6494, 69.7741, 66.0890\n",
    "bev  AP40:69.5448, 62.1173, 57.1881\n",
    "3d   AP40:66.6659, 59.2055, 54.1700\n",
    "aos  AP40:70.00, 64.19, 60.31\n",
    "Pedestrian AP40@0.50, 0.25, 0.25:\n",
    "bbox AP40:75.6494, 69.7741, 66.0890\n",
    "bev  AP40:82.8723, 78.0379, 73.2982\n",
    "3d   AP40:82.6538, 77.1948, 72.8713\n",
    "aos  AP40:70.00, 64.19, 60.31\n",
    "Cyclist AP40@0.50, 0.50, 0.50:\n",
    "bbox AP40:93.8638, 84.2218, 80.1001\n",
    "bev  AP40:92.8451, 75.6214, 71.7649\n",
    "3d   AP40:90.3880, 73.2361, 69.4116\n",
    "aos  AP40:93.71, 83.84, 79.61\n",
    "Cyclist AP40@0.50, 0.25, 0.25:\n",
    "bbox AP40:93.8638, 84.2218, 80.1001\n",
    "bev  AP40:93.9661, 81.6019, 77.2742\n",
    "3d   AP40:93.9661, 81.6019, 77.2742\n",
    "aos  AP40:93.71, 83.84, 79.61\n",
    "Car AP40@0.70, 0.70, 0.70:\n",
    "bbox AP40:97.8348, 94.5482, 94.0081\n",
    "bev  AP40:94.4796, 90.7830, 88.6291\n",
    "3d   AP40:91.8635, 84.5625, 82.4022\n",
    "aos  AP40:97.80, 94.41, 93.80\n",
    "Car AP40@0.70, 0.50, 0.50:\n",
    "bbox AP40:97.8348, 94.5482, 94.0081\n",
    "bev  AP40:97.9316, 96.4609, 94.4074\n",
    "3d   AP40:97.8820, 94.6416, 94.3069\n",
    "aos  AP40:97.80, 94.41, 93.80\n",
    "\n",
    "Overall AP40@easy, moderate, hard:\n",
    "bbox AP40:89.1160, 82.8480, 80.0657\n",
    "bev  AP40:85.6232, 76.1739, 72.5274\n",
    "3d   AP40:82.9724, 72.3347, 68.6612\n",
    "aos  AP40:87.17, 80.81, 77.91\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CenterPoint Voxel 0.1 Circular NMS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the first NuScenes model in the demo. Please remember to update the path of the NuScenes dataset in the mmdetection3d dataset config `configs/_base_/datasets/nus-3d.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "centerpoint_config = {\n",
    "    \"ckpt_before\": \"CenterPoint/centerpoint_01voxel_second_secfpn_circlenms_4x8_cyclic_20e_nus_20220810_030004-9061688e.pth\",\n",
    "    \"ckpt_after\": \"CenterPoint/centerpoint_01voxel_second_secfpn_circlenms_4x8_cyclic_20e_nus_20220810_030004-9061688e.pth\",\n",
    "    \"cfg_path\": \"centerpoint/centerpoint_voxel01_second_secfpn_head-circlenms_8xb4-cyclic-20e_nus-3d.py\",\n",
    "    \"v_spconv\": 1,\n",
    "    \"cfg_options\": \"--cfg-options model.pts_middle_encoder.type=SparseEncoderTS\"\n",
    "}\n",
    "\n",
    "centerpoint_results = mmdet3d_single_demo(centerpoint_config, base_paths, convert=True)\n",
    "print(centerpoint_results.stderr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(centerpoint_results.stdout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(centerpoint_results.stdout)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Expected Outputs: \n",
    "\n",
    "```\n",
    "Evaluating bboxes of pred_instances_3d\n",
    "mAP: 0.5544\n",
    "mATE: 0.2988\n",
    "mASE: 0.2538\n",
    "mAOE: 0.3110\n",
    "mAVE: 0.3039\n",
    "mAAE: 0.1977\n",
    "NDS: 0.6407\n",
    "Eval time: 53.4s\n",
    "\n",
    "Per-class results:\n",
    "Object Class\tAP\tATE\tASE\tAOE\tAVE\tAAE\n",
    "car\t0.845\t0.186\t0.152\t0.113\t0.304\t0.194\n",
    "truck\t0.522\t0.324\t0.185\t0.126\t0.283\t0.240\n",
    "bus\t0.667\t0.354\t0.181\t0.062\t0.535\t0.268\n",
    "trailer\t0.362\t0.546\t0.207\t0.447\t0.208\t0.164\n",
    "construction_vehicle\t0.160\t0.639\t0.414\t0.858\t0.117\t0.334\n",
    "pedestrian\t0.827\t0.165\t0.276\t0.410\t0.244\t0.101\n",
    "motorcycle\t0.529\t0.213\t0.237\t0.292\t0.511\t0.264\n",
    "bicycle\t0.341\t0.169\t0.268\t0.421\t0.229\t0.016\n",
    "traffic_cone\t0.638\t0.162\t0.342\tnan\tnan\tnan\n",
    "barrier\t0.653\t0.230\t0.277\t0.070\tnan\tnan\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torchsparse",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
